# -*- encoding:utf-8

"""
unpack large vpk file, and convert it to psp save data, with a relatively small vpk file.
When there is lnk file in vpk files, rename it or include it in the small vpk.
"""

import zipfile
import os
import os.path
import shutil
import logging
from PyQt4.QtCore import *

logger = logging.Logger('vpk')
logger.setLevel(logging.DEBUG)

class Unpacker(QThread):
    def __init__(self):
        super().__init__()
        self.settask()

    def settask(self, vpk=None, folder=None, rename=True):
        self.task = vpk
        self.folder = folder
        self.rename = rename
        self.progress = 0.0

    def run(self):
        self.unpack()

    def unpack(self):
        if self.task is not None:
            self.unpack_heler()
        self.task = None

    def unpack_heler(self):
        vpk, folder, rename = self.task, self.folder, self.rename
        zfile = zipfile.ZipFile(vpk, 'r')
        if folder is None:
            _, filename = os.path.split(vpk)
            folder = filename[0:-4]
        base = os.path.join(folder, 'app')
        shutil.rmtree(folder, True)
        os.makedirs(base, exist_ok=True)
        smallvpk = zipfile.ZipFile(os.path.join(folder, 'app', 'install.vpk'), 'w')

        length = len(zfile.filelist)
        packlist = []
        lnklist = []

        for i, member in enumerate(zfile.namelist()):
            # extract all files
            zfile.extract(member, base)
            if member.startswith('sce_sys'):
                packlist.append(member)
            if member.endswith('lnk'):
                lnklist.append(member)
            self.progress = i/length

        infolist = ['icon0.png', 'pic0.png', 'param.sfo']
        for info in infolist:
            filename = os.path.join(base,'sce_sys', info)
            shutil.copy(filename, os.path.join(folder, info) )

        for member in packlist + ['eboot.bin']:
            filename = os.path.join(base, member)
            move2vpk(smallvpk, filename, member)
        shutil.rmtree(os.path.join(base,'sce_sys'))


        # lnk file won't be able to transfer to the psv by QCMA
        for member in lnklist:
            lnkfile = os.path.join(base, member)
            if rename:
                os.rename(lnkfile, lnkfile+'.bak')
                logger.debug(lnkfile + ' renamed.')
            else:
                move2vpk(smallvpk, lnkfile, member)

        smallvpk.close()



def unpack(vpk, folder=None, rename=True):
    zfile = zipfile.ZipFile(vpk, 'r')
    if folder is None:
        _, filename = os.path.split(vpk)
        folder = filename[0:-4]
    base = os.path.join(folder, 'app')
    shutil.rmtree(base, True)
    os.makedirs(base, exist_ok=True)
    smallvpk = zipfile.ZipFile(os.path.join(folder, 'app', 'install.vpk'), 'w')

    length = len(zfile.filelist)
    packlist = []
    lnklist = []

    for i, member in enumerate(zfile.namelist()):
        # extract all files
        zfile.extract(member, base)
        if member.startswith('sce_sys'):
            packlist.append(member)
        if member.endswith('lnk'):
            lnklist.append(member)

    infolist = ['icon0.png', 'pic0.png', 'param.sfo']
    for info in infolist:
        filename = os.path.join(base,'sce_sys', info)
        shutil.copy(filename, os.path.join(folder, info) )

    for member in packlist + ['eboot.bin']:
        filename = os.path.join(base, member)
        move2vpk(smallvpk, filename, member)
    shutil.rmtree(os.path.join(base,'sce_sys'))


    # lnk file won't be able to transfer to the psv by QCMA
    for member in lnklist:
        lnkfile = os.path.join(base, member)
        if rename:
            os.rename(lnkfile, lnkfile+'.bak')
            logger.debug(lnkfile + ' renamed.')
        else:
            move2vpk(smallvpk, lnkfile, member)

    smallvpk.close()

def move2vpk(smallvpk, filename, arcname):
    smallvpk.write(filename, arcname)
    logger.debug(arcname + ' added.')
    if os.path.isfile(filename):
        os.remove(filename)
        logger.debug(filename + ' removed.')


if __name__ == '__main__':
    vpk = r'D:\BaiduYunDownload\psv\0233_汉化_PCSG00304_影牢暗黑公主_0102.vpk'
    #vpk = r'D:\BaiduYunDownload\psv\PCSE00449_Rogue_Legacy_FULLGAME_01.00.vpk'
    unpack(vpk, rename=False)